#!/usr/bin/env python3
# coding: utf-8
"""
多 GPU（可选）监控脚本：
  * 每分钟采样一次：温度℃、功耗W、显存占用%、核心利用率%
  * 时间戳使用东八区（Asia/Shanghai）
  * 每 GPU 累计 24 小时（1440 条记录）自动绘图；CSV 与 PNG 独立归档
  * 命令示例：
      单 GPU（默认 GPU0）:      python gpu_monitor.py
      指定多 GPU（如 0 和 1）:   python gpu_monitor.py --gpus 0,1
"""

import argparse
import time
from pathlib import Path
import pandas as pd
import matplotlib.pyplot as plt
import pynvml as nv

# ------------------ 参数 ------------------
SAMPLE_INTERVAL = 60                # 采样周期（秒）
WINDOW_SIZE     = 24 * 60           # 24h * 60min
BASE_DIR        = Path("gpu_monitor_charts")
TZONE           = "Asia/Shanghai"   # 东八区

plt.rcParams["font.size"] = 10
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# ------------------ 工具函数 ------------------
def parse_args():
    p = argparse.ArgumentParser(description="Nvidia GPU 监控")
    p.add_argument("--gpus", default="0",
                   help="要监控的 GPU 索引列表，如 0 或 0,1 (默认 0)")
    return p.parse_args()

def query_stats(handle):
    """返回一次采样结果 dict"""
    ts   = pd.Timestamp.now(tz=TZONE)
    temp = nv.nvmlDeviceGetTemperature(handle, nv.NVML_TEMPERATURE_GPU)
    pwr  = nv.nvmlDeviceGetPowerUsage(handle) / 1000  # mW→W
    mem  = nv.nvmlDeviceGetMemoryInfo(handle)
    mem_pct = round(mem.used / mem.total * 100, 2)
    util = nv.nvmlDeviceGetUtilizationRates(handle).gpu
    return {"timestamp": ts,
            "temperature": temp,
            "power": pwr,
            "memory_pct": mem_pct,
            "util_pct": util}

def save_chart(df: pd.DataFrame, out_dir: Path, seq: int, gpu_idx: int):
    fig, ax = plt.subplots(figsize=(10, 6), dpi=120)
    ax.plot(df["timestamp"], df["temperature"], label="Temperature (℃)")
    ax.plot(df["timestamp"], df["power"],       label="Power (W)")
    ax.plot(df["timestamp"], df["memory_pct"],  label="Memory utilization(%)")
    ax.plot(df["timestamp"], df["util_pct"],    label="Core utilization (%)")
    ax.set_title(f"GPU{gpu_idx} 24h Health status(Graph {seq})")
    ax.set_xlabel("Time (Asia/Shanghai)")
    ax.set_ylabel("Value")
    ax.legend(loc="upper left")
    ax.grid(True, linestyle="--", alpha=0.4)
    fig.autofmt_xdate()
    fname = out_dir / f"gpu{gpu_idx}_24h_{seq:03d}.png"
    fig.savefig(fname, bbox_inches="tight")
    plt.close(fig)
    print(f"[GPU{gpu_idx}] 24h 图已保存：{fname}")

# ------------------ 主逻辑 ------------------
def main():
    args = parse_args()
    gpu_indices = sorted({int(x) for x in args.gpus.split(",") if x.strip().isdigit()})

    nv.nvmlInit()
    try:
        # 初始化每 GPU 的句柄、缓存、序号、输出目录、CSV 路径
        gpus = {}
        for idx in gpu_indices:
            handle = nv.nvmlDeviceGetHandleByIndex(idx)
            out_dir = BASE_DIR / f"gpu{idx}"
            out_dir.mkdir(parents=True, exist_ok=True)
            csv_path = out_dir / f"gpu{idx}_monitor_log.csv"
            buffer   = []
            seq      = 1
            # 若已有历史 CSV，继续累积
            if csv_path.exists():
                df_hist = pd.read_csv(csv_path, parse_dates=["timestamp"])
                buffer  = df_hist.tail(WINDOW_SIZE).to_dict("records")
                seq     = len(df_hist) // WINDOW_SIZE + 1
            gpus[idx] = {"handle": handle,
                         "out_dir": out_dir,
                         "csv": csv_path,
                         "buf": buffer,
                         "seq": seq}

        while True:
            for idx, info in gpus.items():
                try:
                    data = query_stats(info["handle"])
                except nv.NVMLError as e:
                    print(f"[GPU{idx}] 读取失败：{e}")
                    continue

                info["buf"].append(data)
                # 追加一行到 CSV
                pd.DataFrame([data]).to_csv(
                    info["csv"], mode="a", index=False,
                    header=not info["csv"].exists()
                )

                print(f"[{data['timestamp']:%Y-%m-%d %H:%M}] GPU{idx} "
                      f"T={data['temperature']}℃  "
                      f"P={data['power']}W  "
                      f"Mem={data['memory_pct']}%  "
                      f"Util={data['util_pct']}%")

                # 满 24h 绘图
                if len(info["buf"]) >= WINDOW_SIZE:
                    df24 = pd.DataFrame(info["buf"][-WINDOW_SIZE:])
                    save_chart(df24, info["out_dir"], info["seq"], idx)
                    info["seq"] += 1
                    info["buf"].clear()

            # 保证整分采样
            now = time.time()
            time.sleep(max(0, SAMPLE_INTERVAL - now % SAMPLE_INTERVAL))

    finally:
        nv.nvmlShutdown()

if __name__ == "__main__":
    main()
